feat(pt): support shared_dict in linear energy model#5548
Conversation
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds Changesshared_dict parameter sharing for LinearEnergyModel
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
source/tests/pt/model/test_linear_model_shared_dict.py (1)
9-101: 🏗️ Heavy liftAdd coverage for the unexercised shared_dict branches.
These tests only cover descriptor sharing through
get_model. Please add focused regressions for the newupdate_selshared_dict round-trip and a non-descriptor shared link, so the descriptor-restoration logic and the Line 108-Line 144share_paramsbranch cannot regress silently.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/model/test_linear_model_shared_dict.py` around lines 9 - 101, Add two new test methods to the TestLinearEnergySharedDict class to improve coverage. First, create a test that exercises the update_sel shared_dict round-trip functionality to ensure descriptor restoration logic works correctly when sel is modified and restored. Second, create a test that covers a non-descriptor shared link (not just descriptor sharing) to exercise the share_params branch and ensure parameter sharing for non-descriptor shared objects works as expected and cannot regress silently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/model/model/__init__.py`:
- Around line 177-180: The code at line 177 assumes all sub-models in
model_params["models"] have the same type_map but only validates the first one.
Before copying type_map from models[0], iterate through all models in
model_params["models"] to ensure each one has a type_map and that all type_maps
are identical. If any inconsistency is found (missing type_map or differing
values), raise a ValueError with a descriptive message. Only proceed with
copying from models[0] if all models have consistent type_maps.
In `@deepmd/pt/model/model/dp_linear_model.py`:
- Around line 84-99: Replace the assert statements that validate shared_level
ordering and type matching with explicit exception-based validation. Convert
shared_level_base to int using int() function before comparisons, then introduce
a variable to track the previous_shared_level within the loop iterating over
shared_links items. After extracting shared_level_link, validate that it is
greater than or equal to previous_shared_level (not just shared_level_base) and
raise a ValueError if the ordering constraint is violated. Similarly, replace
the assert for class type validation with an explicit exception raise. Apply
these changes consistently to both code paths where this pattern appears (around
the links iteration and the link_class.share_params call).
---
Nitpick comments:
In `@source/tests/pt/model/test_linear_model_shared_dict.py`:
- Around line 9-101: Add two new test methods to the TestLinearEnergySharedDict
class to improve coverage. First, create a test that exercises the update_sel
shared_dict round-trip functionality to ensure descriptor restoration logic
works correctly when sel is modified and restored. Second, create a test that
covers a non-descriptor shared link (not just descriptor sharing) to exercise
the share_params branch and ensure parameter sharing for non-descriptor shared
objects works as expected and cannot regress silently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5c1b8781-2637-4efc-a3cf-a4fe00c95203
📒 Files selected for processing (5)
deepmd/pt/model/model/__init__.pydeepmd/pt/model/model/dp_linear_model.pydeepmd/pt/utils/multi_task.pydeepmd/utils/argcheck.pysource/tests/pt/model/test_linear_model_shared_dict.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5548 +/- ##
==========================================
- Coverage 82.23% 82.20% -0.03%
==========================================
Files 894 898 +4
Lines 102002 103689 +1687
Branches 4276 4432 +156
==========================================
+ Hits 83877 85236 +1359
- Misses 16823 17059 +236
- Partials 1302 1394 +92 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Share linear model parameters in resume mode during construction so descriptor statistics are not required before training computes them. Also validate linear shared links and inferred type maps explicitly. Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/model/model/__init__.py`:
- Around line 185-190: The ValueError raised when sub_model_params type_map
values don't match is missing the model index information, making it difficult
to identify which sub-model has the mismatch. Include the idx variable (which is
already captured in the loop) in the error message string to clearly indicate
which sub-model has the mismatched type_map value. Update the error message in
the ValueError to include the model index along with the existing description.
🪄 Autofix (Beta)
❌ Autofix failed (check again to retry)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: cd4f9bab-18aa-49cb-80c2-5f5a57d58352
📒 Files selected for processing (2)
deepmd/pt/model/model/__init__.pydeepmd/pt/model/model/dp_linear_model.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/model/dp_linear_model.py
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. An unexpected error occurred while generating fixes: Not Found - https://docs.github.com/rest/git/refs#get-a-reference |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt/model/test_linear_model_shared_dict.py (1)
18-18: ⚡ Quick winConsider a more explicit assertion.
Using
assertTrueon a dict checks if it's non-empty, but this intent is not immediately clear. Consider using a more explicit assertion.✨ More explicit alternative
- self.assertTrue(descriptor0.se_atten._modules) + self.assertGreater(len(descriptor0.se_atten._modules), 0, "Expected se_atten to have modules")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/model/test_linear_model_shared_dict.py` at line 18, Replace the implicit truthiness check on the dictionary descriptor0.se_atten._modules with a more explicit assertion that clearly communicates the intent to verify the dictionary is non-empty. Instead of using assertTrue on the dict directly, use an assertion method like assertGreater with len() to explicitly check the dictionary length is greater than zero, making the intent immediately clear to readers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/pt/model/test_linear_model_shared_dict.py`:
- Line 18: Replace the implicit truthiness check on the dictionary
descriptor0.se_atten._modules with a more explicit assertion that clearly
communicates the intent to verify the dictionary is non-empty. Instead of using
assertTrue on the dict directly, use an assertion method like assertGreater with
len() to explicitly check the dictionary length is greater than zero, making the
intent immediately clear to readers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 87b86aff-8314-4e17-aeb1-4f0e2b935577
📒 Files selected for processing (2)
deepmd/pt/model/model/__init__.pysource/tests/pt/model/test_linear_model_shared_dict.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/model/init.py
| if "hybrid" in shared_type: | ||
| hybrid_index = int(shared_type.split("_")[-1]) | ||
| return sub_model.descriptor.descriptor_list[hybrid_index] |
There was a problem hiding this comment.
The hybrid-descriptor branch here (descriptor_list[hybrid_index]) is reachable but never exercised. The only test that builds a hybrid descriptor (test_shared_dict_update_sel_round_trip) drives update_sel, not share_params; the three tests that actually call share_params (via get_model) all use a plain dpa1 descriptor and only hit the shared_type == "descriptor" path above. CLAUDE.md: "When adding a new feature or API, provide tests that exercise every reachable code path." Consider a test that shares one component of a hybrid descriptor between two linear sub-models. Non-blocking.
| if isinstance(descriptor_ref, str): | ||
| ret_jdata["shared_dict"][get_shared_key(descriptor_ref)] = ( | ||
| updated_sub_model["descriptor"] | ||
| ) | ||
| elif ( | ||
| isinstance(descriptor_ref, dict) | ||
| and descriptor_ref.get("type") == "hybrid" | ||
| ): | ||
| updated_descriptor = updated_sub_model["descriptor"] | ||
| for hybrid_idx, hybrid_ref in enumerate(descriptor_ref["list"]): | ||
| if isinstance(hybrid_ref, str): | ||
| ret_jdata["shared_dict"][get_shared_key(hybrid_ref)] = ( | ||
| updated_descriptor["list"][hybrid_idx] | ||
| ) | ||
| else: | ||
| ret_jdata["models"][idx]["descriptor"]["list"][hybrid_idx] = ( | ||
| updated_descriptor["list"][hybrid_idx] | ||
| ) | ||
| else: | ||
| ret_jdata["models"][idx]["descriptor"] = updated_sub_model["descriptor"] |
There was a problem hiding this comment.
The descriptor write-back has three branches: a string ref (isinstance(descriptor_ref, str), L320), a hybrid dict (L324), and an inline dict (the final else, L339). test_shared_dict_update_sel_round_trip only uses a hybrid descriptor, so only the middle branch runs; the string-ref and inline-dict branches are untested. CLAUDE.md: "UTs should cover all code paths, including both branches of boolean conditions." Non-blocking.
| if "type_map" not in sub_model_params: | ||
| raise ValueError( | ||
| f"Linear sub-model {idx} must define type_map when " | ||
| "linear_ener has no top-level type_map." | ||
| ) | ||
| first_type_map = model_params["models"][0]["type_map"] | ||
| for idx, sub_model_params in enumerate(model_params["models"][1:], start=1): | ||
| if sub_model_params["type_map"] != first_type_map: | ||
| raise ValueError( | ||
| f"Linear sub-model {idx} type_map differs from sub-model 0. " | ||
| "All type_map values must be identical when linear_ener " |
There was a problem hiding this comment.
The new type_map validation raises in two cases — a sub-model missing type_map (L180-182) and a sub-model whose type_map differs from sub-model 0 (L186-189) — but neither raise is covered. The tests only hit the happy path where all sub-models agree. CLAUDE.md: "When adding a new feature or API, provide tests that exercise every reachable code path." Consider two assertRaises cases. Non-blocking.
3 participants
Summary
shared_dictsupport to the PyTorchlinear_enermodel config, reusing the multi-task shared-parameter preprocessor.update_sel, including configs wheretype_maponly lives inshared_dict.type_mapplus shared descriptors.Closes #4412.
Supersedes/continues #4933.
Tests
uvx ruff check deepmd/pt/model/model/dp_linear_model.py deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py deepmd/utils/argcheck.py source/tests/pt/model/test_linear_model_shared_dict.pyuvx ruff format deepmd/pt/model/model/dp_linear_model.py deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py deepmd/utils/argcheck.py source/tests/pt/model/test_linear_model_shared_dict.pypython3 -m py_compile deepmd/pt/model/model/dp_linear_model.py deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py deepmd/utils/argcheck.py source/tests/pt/model/test_linear_model_shared_dict.pyNote: local
pytest -q source/tests/pt/model/test_linear_model_shared_dict.pycould not run in this worktree because the compileddeepmd.libextension is unavailable in the local environment.Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
Summary by CodeRabbit
shared_dictsupport forlinear_enermodels, enabling parameter sharing across atomic sub-models (including descriptor and fitting-network sharing).linear_enerselection updates to round-tripshared_dictconfigurations while preserving shared descriptor references.type_mapresolution under shared configurations (enforcing consistent sub-modeltype_mapwhen needed).type_map,shared_linkspopulation, hybrid descriptors, andupdate_selround-trip behavior.